from langchain.prompts import PromptTemplate
from pydantic import BaseModel, Field
from datetime import datetime
import json
from LLM_Prompt.siliconflow import SiliconFlow

# 使用 SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()  # 替换成你实际的 SiliconFlow 模型接口

classification_prompt = PromptTemplate(
    input_variables=["question", "current_date"],
    template="""  
# 任务说明

你负责根据用户输入的问题，从下方 API 中选择最相关的一个，并准确提取参数。

## 可用 API
1. `vessel_trajectory`:
   - **用途**：获取船只的航线轨迹数据。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

2. `checkpoint_passage`:
   - **用途**：监控船只经过卡口或检查点的情况。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

3. `violation_detection`:
   - **用途**：查询船只的违规行为或预警。
   - **必需参数**：
     - `boatName`（船只名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

4. `checkpoint_traffic`:
   - **用途**：查询特定卡口通过的船舶信息。
   - **必需参数**：
     - `checkpointName`（卡口名称）
     - `startTime`（查询起始时间，ISO8601 格式）
     - `endTime`（查询结束时间，ISO8601 格式）

5. `current_waterway_status`:
   - **用途**：查询特定航道或港区当前的船舶情况。
   - **必需参数**：
     - `waterwayName`（航道名称）

6. `basic_information`:
   - **用途**：查询某条船的基本信息。
   - **必需参数**：
     - `boatName`（船只名称）

## 特殊要求

1. **分类**：对用户输入的问题进行意图分类，只能选择一个最相关 API。
2. **参数提取**：
    - boatName：提取船名（如有）。
    - startTime / endTime：需为 ISO8601 格式，例如 `2025-04-09 00:00:00`。如未指定时间则返回 "未知"。遇“今天”、“昨天”、“前天”等请按当前日期 `{current_date}` 计算。
    - checkpointName：如涉及卡口，请（模糊）匹配如下：杭申线卡口、长湖申线卡口、金汇港(北闸)、油墩港(油墩港船闸)、苏申内港线(上港宜东)、吴淞江(泗江口站)、黄浦江(松浦大桥南岸)、龙泉港(叶榭塘船闸)、油墩港(东大盈船闸)、大治河船闸、苏申外港线省际检查站、杨思船闸口、东沟船闸口、苏申内港线省际检查站、平申线省际检查站
    - waterwayName：如有，提取航道名。
3. **统计性问题（如总数、统计、汇总等）统一返回 `not matched` API。**
4. **只返回一个最合适的 API 及其参数**

## 输出格式

仅返回以下 JSON 结构：

```json
{{  
  "api": "API 名称",  
  "parameters": {{  
    "boatName": "示例船只名称或 null",  
    "startTime": "2025-04-08 00:00:00 或 null",  
    "endTime": "2025-04-08 23:59:59 或 null",  
    "checkpointName": "卡口名称或 null",  
    "waterwayName": "航道名称或 null"  
  }}  
}} 
```

### 示例  
输入  
问题：{question}  

输出  
严格返回上述 JSON 格式内容。

### 说明

- 请勿输出除 JSON 外的任何内容。
- 如匹配多个 API，仅选择并返回你认为最合适的一个。
- 未提及参数填 null。  
- 时间相关按规则转换，无则填 null。
"""
)

# Step 4: Pydantic 数据验证器
class ExtractedData(BaseModel):
    boat_name: str = Field(..., description="船的名字")
    date: str = Field(..., description="日期，格式为 YYYY-MM-DD")

    @staticmethod
    def validate_output(response: str):
        try:
            data = ExtractedData.model_validate_json(response)  # 替换 parse_raw
            return data.model_dump()  # 替换 dict
        except Exception as e:
            raise ValueError(f"数据解析错误: {e}")


# 清理模型输出工具
def clean_model_output(raw_output: str) -> str:
    """
    清理大模型的输出，去除额外的 Markdown 等格式符号，确保输出是合法 JSON。
    """
    # 去掉 Markdown 格式的代码块标记（如 ```json ... ```)
    if raw_output.startswith("```json"):
        raw_output = raw_output.strip("```json").strip("`\n")
    elif raw_output.startswith("```"):
        raw_output = raw_output.strip("```").strip("`\n")

    try:
        # 确保最终输出为合法的 JSON 格式（如有多余的空格，再清理一下）
        json.loads(raw_output)  # 测试是否是有效 JSON
    except json.JSONDecodeError as e:
        raise ValueError(f"清理后的输出仍不是有效 JSON: {e}")

    return raw_output



# Step 5: 链式操作
classification_chain = classification_prompt | llm


def classify_and_extract(question):
    # 动态获取当前日期
    current_date = datetime.now().strftime("%Y-%m-%d")

    # 1. 问题分类 提取关键信息
    classification_extract_json = classification_chain.invoke({"question": question, "current_date":current_date})

    # 清理模型输出
    cleaned_result = clean_model_output(classification_extract_json)

    return cleaned_result


# 测试代码
test_queries = [
        # "莲花号前天的航道轨迹",
        # "莲花号2025年1月20号的航道轨迹",
        # "船舶 ABC123 2025年2月20日的行驶轨迹？",
        # "船舶 XYZ789 昨天经过那些卡口？",
        # "船舶 DEF456 本周是否有违规行为（预警）？",
        # "船舶 GHI789 上周去过哪些码头？去过哪些区域？途径那些船闸？途径那些航道？",
        # "卡口 南京长江大桥 2025年2月20日到22日有哪些船舶经过？有哪些违规船舶？有哪些集装箱船舶经过等",
        # "航道 长江主航道 当前有哪些船舶？有哪些违规船舶？有哪些类别船舶？",
        # "2025年2月20日至22日集装箱船的行驶轨迹？",
        # "昨天有哪些集装箱船经过南京长江大桥卡口？",
        # "本周有没有集装箱船存在违规行为？",
        # "长江主航道当前有哪些集装箱船？",
        # "上海2025年1月10日天气怎么样",
        # "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        # "江海货368 2025年2月20日航行记录。",
        # "船舶 江海货1819 2024年10月1日到11月30号经过那些卡口？有哪些违规的行为？",
        # "分析2025年1月油墩港(东大盈船闸)每天的流量组成",
        # "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
        # "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。"
]

# for example_question in test_queries:
#   print('example_question:', example_question)
#   result = classify_and_extract(example_question)
#   print(result)

# question = "船舶 苏无锡货89898 是什么类型船舶？"
# result = classify_and_extract(question)
# print(result)